# -*- coding:utf-8 -*-

# @Time    : 2023/5/12 16:57
# @Author  : lishuaichao
# @Email   : lishuaichao@lingxi.ai
# @File    : faq.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import json
import logging
import os
import re
import pandas as pd

from langchain import FAISS, OpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain.schema import Document
import platform

logging.basicConfig(format='%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s',
                    level=logging.INFO)
os.environ["OPENAI_API_KEY"] = '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'


class Faq:
    def __init__(self):
        self.user_job = 'user_job'
        self.solution = 'solution'
        sys_platform = platform.platform().lower()
        print(sys_platform)
        if 'macos' in sys_platform:
            self.folder_path = 'objection_handling/'
        else:
            self.folder_path = 'bot/insurance_sales_gpt/objection_handling/'
        self.embeddings = OpenAIEmbeddings()
        self.llm = OpenAI(temperature=0.9)

    def get_db_name(self, conversation_stage: int):
        if conversation_stage < 4:
            return self.user_job
        else:
            return self.solution


def get_extension_questions(standard_q: str):
    df = pd.read_excel('knowledge_base/FAQ-341-黑牛.xlsx')
    extension_questions = df.to_dict('records')
    # 获取标准问相同的扩展问
    extension_questions = [extension_question['扩展问'] for extension_question in extension_questions
                           if standard_q == extension_question.get('标准问')]
    return extension_questions


class FaqSave(Faq):
    def __init__(self):
        super().__init__()

    def save_from_csv(self, csv_path, conversation_stage: int):
        """
        保存FAQ到向量库
        Args:
            conversation_stage: 阶段
            csv_path: csv文件路径, 列包含 标准问 扩展问 答案

        Returns: 无

        """
        df = pd.read_csv(csv_path)
        faqs = df.to_dict('records')
        logging.info(f'{csv_path}加载成功, 共有{len(faqs)}条数据')
        docs = []
        for faq in faqs:
            questions = get_extension_questions(faq.get('标准问'))
            if len(questions) > 10:
                questions = questions[:10]
            logging.info(f"标准问:{faq.get('标准问')}, 扩展问:{questions}")
            docs.append(Document(page_content=f"标准问:{faq.get('标准问')}\n扩展问:{';'.join(questions)}\n", metadata={'答案': faq.get('答案')}))
        db = FAISS.from_documents(docs, self.embeddings)
        db.save_local(folder_path=self.folder_path, index_name=self.get_db_name(conversation_stage))


class FaqQuery(Faq):
    def __init__(self):
        super().__init__()
        self.db_dict = {
            self.user_job: FAISS.load_local(folder_path=self.folder_path, index_name=self.user_job, embeddings=self.embeddings),
            self.solution: FAISS.load_local(folder_path=self.folder_path, index_name=self.solution, embeddings=self.embeddings)
        }
        self.prompt = (
            "你是一名百万医疗保险的电话销售坐席,你的任务是通过电话劝服用户投保百万医疗保险,"
            "坐席跟用户聊天内容的上下文是:\n"
            "===\n{context}\n"
            "===\n"
            "下面是几个FAQ\n"
            "---\n"
            "{context_str}\n"
            "---\n"
            "用户的最后一句话是:{query_str}\n"
            "请你充分理解对话,如果用户有疑问请判断FAQ的答案是否能解决用户疑问,如果用户拒绝坐席请判断FAQ能否劝服用户\n"
            "如果能解决或劝服请输出FAQ编号的内容,如果所有FAQ都不能解决用户问题或异议请输出'无',如果用户没有明确问题或异议请输出'无'\n"
            "### Response:\n"
        )

    def query(self, query_str: str, chat_context: str = None, conversation_stage: int = -1) -> (str, str, bool):
        """
        FAQ查询函数
        Args:
            query_str: 用户问题
            chat_context: 聊天的上下文
            conversation_stage: 服务阶段
        Returns: FAQ答案

        """
        if query_str is None or query_str == '':
            return None, None, False
        chat_context = self.segmenting_chat_context(chat_context)
        docs = self.db_dict.get(self.get_db_name(conversation_stage)).similarity_search(query_str, k=3)
        faqs = []
        for index, doc in enumerate(docs):
            logging.info(doc)
            # 正则提取出标准问
            # question = re.search('标准问:(.*?)\n扩展问', doc.page_content).group(1)
            # faqs.append(f"FAQ编号:{index}\n标准问:{question}\n答案:{doc.metadata.get('答案')}")
            faqs.append(f"FAQ编号:{index}\n{doc.page_content}答案:{doc.metadata.get('答案')}")
        context_str = '\n---\n'.join(faqs)
        prompt = self.prompt.format(query_str=query_str, context_str=context_str, context=chat_context)
        logging.info(f'prompt:{prompt}')
        response = self.llm(prompt)
        logging.info(f'用户问题:{query_str}')
        logging.info(f'大模型返回的编号:{response}')
        if response.startswith('无'):
            return None, None, False
        else:
            number = re.search('\d+', response).group()
            if 0 <= int(number) < len(docs):
                answer = docs[int(number)].metadata.get('答案')
                # 正则提取出标准问
                question = re.search('标准问:(.*?)\n扩展问', docs[int(number)].page_content).group(1)
                logging.info(f'FAQ:{docs[int(number)].page_content}答案:{answer}')
                return question, answer, '再见' not in answer
            else:
                return None, None, False

    @staticmethod
    def segmenting_chat_context(chat_context):
        if chat_context is not None:
            lines = chat_context.splitlines()
            if len(lines) >= 6:
                chat_context = '\n'.join(lines[-6:])
        return chat_context


if __name__ == '__main__':
    # faq_save = FaqSave()
    # faq_save.save('knowledge_base/faq_list.json')
    # faq_save.save_from_csv('knowledge_base/user_job.csv', 1)
    # faq_save.save_from_csv('knowledge_base/solution.csv', 10)
    faq_query = FaqQuery()
    context = ''
    # print(faq_query.query('对,请问你是谁啊?', chat_context=context, conversation_stage=1))
    # 打开txt文件并读所有行
    with open('knowledge_base/成龙--NBS风险激发与案例激发.trans.txt', 'r', encoding='utf-8') as f:
        lines = f.readlines()
        for line in lines:
            context += line
            if line.startswith('用户:'):
                print(faq_query.query(line, chat_context=context, conversation_stage=1))
                print(context)
                print('-----------------------------')
